# pylint: disable=line-too-long,useless-suppression
# ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------

"""
DESCRIPTION:
    This sample demonstrates how to interact with the memory store to add and retrieve memory
    using the synchronous AIProjectClient. It uses some additional operations compared
    to the basic memory sample.

    See also /samples/agents/tools/sample_agent_memory_search.py that shows
    how to use the Memory Search Tool in a prompt agent.

USAGE:
    python sample_memory_advanced.py

    Before running the sample:

    pip install "azure-ai-projects>=2.0.0b1" python-dotenv

    Deploy a chat model (e.g. gpt-4.1) and an embedding model (e.g. text-embedding-3-small).
    Once you have deployed models, set the deployment name in the variables below.

    Set these environment variables with your own values:
    1) AZURE_AI_PROJECT_ENDPOINT - The Azure AI Project endpoint, as found in the Overview
       page of your Microsoft Foundry portal.
    2) AZURE_AI_CHAT_MODEL_DEPLOYMENT_NAME - The deployment name of the chat model, as found under the "Name" column in
       the "Models + endpoints" tab in your Microsoft Foundry project.
    3) AZURE_AI_EMBEDDING_MODEL_DEPLOYMENT_NAME - The deployment name of the embedding model, as found under the
       "Name" column in the "Models + endpoints" tab in your Microsoft Foundry project.
"""

import os
from dotenv import load_dotenv
from azure.core.exceptions import ResourceNotFoundError
from azure.identity import DefaultAzureCredential
from azure.ai.projects import AIProjectClient
from azure.ai.projects.models import (
    MemoryStoreDefaultDefinition,
    MemoryStoreDefaultOptions,
    MemorySearchOptions,
    ResponsesUserMessageItemParam,
    ResponsesAssistantMessageItemParam,
)

load_dotenv()

endpoint = os.environ["AZURE_AI_PROJECT_ENDPOINT"]

with (
    DefaultAzureCredential(exclude_interactive_browser_credential=False) as credential,
    AIProjectClient(endpoint=endpoint, credential=credential) as project_client,
):

    # Delete memory store, if it already exists
    memory_store_name = "my_memory_store"
    try:
        project_client.memory_stores.delete(memory_store_name)
        print(f"Memory store `{memory_store_name}` deleted")
    except ResourceNotFoundError:
        pass

    # Create memory store with advanced options
    options = MemoryStoreDefaultOptions(
        user_profile_enabled=True,
        user_profile_details="Preferences and interests relevant to coffee expert agent",
        chat_summary_enabled=True,
    )
    definition = MemoryStoreDefaultDefinition(
        chat_model=os.environ["AZURE_AI_CHAT_MODEL_DEPLOYMENT_NAME"],
        embedding_model=os.environ["AZURE_AI_EMBEDDING_MODEL_DEPLOYMENT_NAME"],
        options=options,
    )
    memory_store = project_client.memory_stores.create(
        name=memory_store_name,
        description="Example memory store for conversations",
        definition=definition,
    )
    print(f"Created memory store: {memory_store.name} ({memory_store.id}): {memory_store.description}")

    # Set scope to associate the memories with.
    # You can also use "{{$userId}}" to take the oid of the request authentication header.
    scope = "user_123"

    # Extract memories from messages and add them to the memory store
    user_message = ResponsesUserMessageItemParam(
        content="I prefer dark roast coffee and usually drink it in the morning"
    )
    update_poller = project_client.memory_stores.begin_update_memories(
        name=memory_store.name,
        scope=scope,
        items=[user_message],  # Pass conversation items that you want to add to memory
        update_delay=300,  # Keep default inactivity delay before starting update
    )
    print(f"Scheduled memory update operation (Update ID: {update_poller.update_id}, Status: {update_poller.status()})")

    # Extend the previous update with another update and more messages
    new_message = ResponsesUserMessageItemParam(content="I also like cappuccinos in the afternoon")
    new_update_poller = project_client.memory_stores.begin_update_memories(
        name=memory_store.name,
        scope=scope,
        items=[new_message],
        previous_update_id=update_poller.update_id,  # Extend from previous update ID
        update_delay=0,  # Trigger update immediately without waiting for inactivity
    )
    print(
        f"Scheduled memory update operation (Update ID: {new_update_poller.update_id}, Status: {new_update_poller.status()})"
    )

    # As first update has not started yet, the new update will cancel the first update and cover both sets of messages
    print(
        f"Superseded first memory update operation (Update ID: {update_poller.update_id}, Status: {update_poller.status()})"
    )

    new_update_result = new_update_poller.result()
    print(
        f"Second update {new_update_poller.update_id} completed with {len(new_update_result.memory_operations)} memory operations"
    )
    for operation in new_update_result.memory_operations:
        print(
            f"  - Operation: {operation.kind}, Memory ID: {operation.memory_item.memory_id}, Content: {operation.memory_item.content}"
        )

    # Retrieve memories from the memory store
    query_message = ResponsesUserMessageItemParam(content="What are my morning coffee preferences?")
    search_response = project_client.memory_stores.search_memories(
        name=memory_store.name, scope=scope, items=[query_message], options=MemorySearchOptions(max_memories=5)
    )
    print(f"Found {len(search_response.memories)} memories")
    for memory in search_response.memories:
        print(f"  - Memory ID: {memory.memory_item.memory_id}, Content: {memory.memory_item.content}")

    # Perform another search using the previous search as context
    agent_message = ResponsesAssistantMessageItemParam(
        content="You previously indicated a preference for dark roast coffee in the morning."
    )
    followup_query = ResponsesUserMessageItemParam(
        content="What about afternoon?"  # Follow-up assuming context from previous messages
    )
    followup_search_response = project_client.memory_stores.search_memories(
        name=memory_store.name,
        scope=scope,
        items=[agent_message, followup_query],
        previous_search_id=search_response.search_id,
        options=MemorySearchOptions(max_memories=5),
    )
    print(f"Found {len(followup_search_response.memories)} memories")
    for memory in followup_search_response.memories:
        print(f"  - Memory ID: {memory.memory_item.memory_id}, Content: {memory.memory_item.content}")

    # Delete memories for the current scope
    project_client.memory_stores.delete_scope(name=memory_store.name, scope=scope)
    print(f"Deleted memories for scope '{scope}'")

    # Delete memory store
    project_client.memory_stores.delete(memory_store.name)
    print(f"Deleted memory store `{memory_store.name}`")
